import os
from typing import Optional, List, Dict
import httpx
from mcp.server.fastmcp import FastMCP


def get_api_key() -> str:
    """Get API key from environment variables"""
    api_key = os.getenv("HEFENG_API_KEY")
    if not api_key:
        raise ValueError("HEFENG_API_KEY环境变量未设置，请检查客户端配置")
    return api_key


def format_weather(weather_data: Dict) -> str:
    """
    格式化天气数据

    参数:
        weather_data: 天气原始数据

    返回:
        格式化后的天气数据字符串
    """
    formatted = ""
    if "now" in weather_data:
        # 实时天气数据格式化
        now_data = weather_data["now"]
        formatted = f"当前天气：{now_data['text']}，温度{now_data['temp']}℃，体感温度{now_data['feelsLike']}℃，湿度{now_data['humidity']}%，风速{now_data['windSpeed']}km/h"
    elif "hourly" in weather_data:
        # 小时预报数据格式化
        for hourly in weather_data["hourly"]:
            formatted += f"{hourly['fxTime'][11:16]}: {hourly['text']}, 温度{hourly['temp']}℃, 降水概率{hourly['pop']}%\n"
    elif "daily" in weather_data:
        # 日预报数据格式化
        for daily in weather_data["daily"]:
            formatted += f"{daily['fxDate']}: {daily['textDay']}, 温度{daily['tempMin']}~{daily['tempMax']}℃\n"
    return formatted


mcp = FastMCP("QWeather MCP Server")


@mcp.tool()
async def get_weather(city: str, days: str) -> str:
    """
    获取城市天气预报（格式化输出）

    参数:
        city: 城市ID，例如："101010100"
        days: 预报类型，可选值：
              "now" - 实时天气
              "24h"/"72h"/"168h" - 小时天气预报
              "3d"/"7d"/"10d"/"15d"/"30d" - 天气预报

    返回:
        格式化后的天气预报字符串

    异常:
        ValueError: 参数错误或API响应格式错误
        httpx.RequestError: 网络请求失败
    """
    # 验证参数
    valid_days = ["now", "24h", "72h", "168h", "3d", "7d", "10d", "15d", "30d"]
    if days not in valid_days:
        raise ValueError(f"Invalid days parameter. Must be one of: {valid_days}")

    HEFENG_WEATHER_URL = (
        "https://nb6yvwk53a.re.qweatherapi.com/v7/weather/{forecast_type}"
    )

    location_info = await get_city_coordinates(city)
    location_id = location_info["id"]
    api_key = get_api_key()
    headers = {"X-QW-Api-Key": api_key}
    params = {"location": location_id}

    async with httpx.AsyncClient() as client:
        try:
            response = await client.get(
                HEFENG_WEATHER_URL.format(forecast_type=days),
                headers=headers,
                params=params,
            )
            response.raise_for_status()
            data = response.json()

            if data.get("code") != "200":
                raise ValueError(
                    f"Weather API Error: {data.get('message', 'Unknown error')}"
                )
            return format_weather(data)
        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


@mcp.tool()
async def get_weather_warning(city: str) -> str:
    """
    获取城市天气灾害预警（格式化输出）

    参数:
        city: 城市ID，例如："101010100"

    返回:
        格式化后的天气灾害预警字符串，包含预警标题、开始时间、结束时间和预警详情

    异常:
        ValueError: API响应格式错误
        httpx.RequestError: 网络请求失败
    """
    HEFENG_WARNING_URL = "https://nb6yvwk53a.re.qweatherapi.com/v7/warning/now"

    location_info = await get_city_coordinates(city)
    location_id = location_info["id"]
    api_key = get_api_key()
    headers = {"X-QW-Api-Key": api_key}
    params = {"location": location_id}

    async with httpx.AsyncClient() as client:
        try:
            response = await client.get(
                HEFENG_WARNING_URL, headers=headers, params=params
            )
            response.raise_for_status()
            data = response.json()

            if data.get("code") != "200":
                raise ValueError(
                    f"Weather Warning API Error: {data.get('message', 'Unknown error')}"
                )

            # 格式化预警数据
            if not data.get("warning"):
                return "当前无天气预警信息"

            formatted_warnings = []
            for warning in data["warning"]:
                formatted_warning = f"预警标题: {warning.get('title', 'N/A')}\n"
                formatted_warning += f"开始时间: {warning.get('startTime', 'N/A')}\n"
                formatted_warning += f"结束时间: {warning.get('endTime', 'N/A')}\n"
                formatted_warning += f"预警详情: {warning.get('text', 'N/A')}\n"
                formatted_warning += "-" * 30 + "\n"
                formatted_warnings.append(formatted_warning)

            return "".join(formatted_warnings)
        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


@mcp.tool()
async def get_city_coordinates(city: str) -> dict:
    """
    获取城市经纬度坐标

    参数:
        city: 城市名称，例如："北京"、"上海"

    返回:
        包含城市经纬度信息的字典，包括:
        - name: 城市名称
        - id: 城市ID
        - lat: 纬度
        - lon: 经度
        - adm1: 省级行政区
        - adm2: 市级行政区
        - country: 国家

    异常:
        ValueError: API响应格式错误
        httpx.RequestError: 网络请求失败
    """
    HEFENG_GEO_URL = "https://nb6yvwk53a.re.qweatherapi.com/geo/v2/city/lookup"

    api_key = get_api_key()
    headers = {"X-QW-Api-Key": api_key}
    params = {"location": city}

    async with httpx.AsyncClient() as client:
        try:
            response = await client.get(HEFENG_GEO_URL, headers=headers, params=params)
            response.raise_for_status()
            data = response.json()

            if data.get("code") != "200":
                raise ValueError(
                    f"GeoAPI Error: {data.get('message', 'Unknown error')}"
                )

            location = data["location"][0]
            return {
                "name": location["name"],
                "id": location["id"],
                "lat": location["lat"],
                "lon": location["lon"],
                "adm1": location["adm1"],
                "adm2": location["adm2"],
                "country": location["country"],
            }
        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


@mcp.tool()
async def get_air_quality(city: str, forecast_type: str = "realtime") -> str:
    """
    获取城市空气质量数据（实时、小时或每日预报）

    参数:
        city: 城市名称，例如："北京"、"上海"
        forecast_type: 预报类型，'realtime'表示实时，'hourly'表示小时预报，'daily'表示每日预报，默认为'realtime'

    返回:
        格式化后的空气质量数据字符串

    异常:
        ValueError: 参数错误或API响应格式错误
        httpx.RequestError: 网络请求失败
    """
    # 首先获取城市的经纬度坐标
    city_info = await get_city_coordinates(city)
    latitude = city_info["lat"]
    longitude = city_info["lon"]

    # 根据预报类型选择API端点
    if forecast_type == "realtime":
        HEFENG_AIR_URL = (
            "https://nb6yvwk53a.re.qweatherapi.com/airquality/v1/current/{lat}/{lon}"
        )
    elif forecast_type == "hourly":
        HEFENG_AIR_URL = (
            "https://nb6yvwk53a.re.qweatherapi.com/airquality/v1/hourly/{lat}/{lon}"
        )
    elif forecast_type == "daily":
        HEFENG_AIR_URL = (
            "https://nb6yvwk53a.re.qweatherapi.com/airquality/v1/daily/{lat}/{lon}"
        )
    else:
        raise ValueError("无效的预报类型，支持 'realtime', 'hourly' 或 'daily'")

    api_key = get_api_key()
    headers = {"X-QW-Api-Key": api_key}

    url = HEFENG_AIR_URL.format(lat=latitude, lon=longitude)

    async with httpx.AsyncClient() as client:
        try:
            response = await client.get(url, headers=headers)
            response.raise_for_status()
            data = response.json()

            if forecast_type == "realtime":
                if "indexes" not in data:
                    raise ValueError("Air Quality API Error: No data returned")

                # 编码处理：优先使用US EPA标准，其次使用QAQI，最后使用第一个指数
                primary_index = None
                preferred_codes = ["us-epa", "qaqi"]

                # 尝试按优先级选择指数
                for code in preferred_codes:
                    for index in data["indexes"]:
                        if index["code"] == code:
                            primary_index = index
                            break
                    if primary_index:
                        break

                # 如果没有找到优先指数，使用第一个
                if not primary_index:
                    primary_index = data["indexes"][0]

                # 提取主要污染物数据（可能为空）
                primary_pollutant = primary_index.get("primaryPollutant")
                if not primary_pollutant:
                    primary_pollutant = {
                        "code": "none",
                        "name": "无",
                        "fullName": "无首要污染物",
                    }

                # 提取各污染物浓度
                pollutants = {}
                if "pollutants" in data:
                    for pollutant in data["pollutants"]:
                        code = pollutant["code"]
                        pollutants[code] = {
                            "name": pollutant["name"],
                            "fullName": pollutant["fullName"],
                            "value": pollutant["concentration"]["value"],
                            "unit": pollutant["concentration"]["unit"],
                        }

                # 提取健康建议（如果存在）
                health_effect = None
                health_advice_general = None
                health_advice_sensitive = None

                if "health" in primary_index:
                    health_info = primary_index["health"]
                    health_effect = health_info.get("effect")
                    if "advice" in health_info:
                        advice = health_info["advice"]
                        health_advice_general = advice.get("generalPopulation")
                        health_advice_sensitive = advice.get("sensitivePopulation")

                # 提取颜色信息（如果存在）
                color_info = None
                if "color" in primary_index:
                    color = primary_index["color"]
                    color_info = {
                        "red": color.get("red"),
                        "green": color.get("green"),
                        "blue": color.get("blue"),
                        "alpha": color.get("alpha"),
                    }

                air_quality_info = {
                    "city": city_info["name"],
                    "index_code": primary_index["code"],
                    "index_name": primary_index["name"],
                    "aqi": primary_index["aqi"],
                    "aqiDisplay": primary_index["aqiDisplay"],
                    "level": primary_index.get("level"),
                    "category": primary_index["category"],
                    "color": color_info,
                    "primaryPollutant": {
                        "code": primary_pollutant["code"],
                        "name": primary_pollutant["name"],
                        "fullName": primary_pollutant["fullName"],
                    },
                    "pollutants": pollutants,
                    "health_effect": health_effect,
                    "health_advice_general": health_advice_general,
                    "health_advice_sensitive": health_advice_sensitive,
                    "type": "realtime",
                }

                return format_air_quality(air_quality_info)

            elif forecast_type == "hourly":
                if "hours" not in data:
                    raise ValueError("小时预报API错误：未返回数据")

                # 处理小时预报数据
                hourly_forecast = []
                for hour_data in data["hours"]:
                    # 提取时间（只保留小时:分钟）
                    forecast_time = hour_data["forecastTime"][11:16]

                    # 查找首选指数 (QAQI)
                    qaqi_index = next(
                        (idx for idx in hour_data["indexes"] if idx["code"] == "qaqi"),
                        None,
                    )

                    if not qaqi_index:
                        qaqi_index = hour_data["indexes"][0]

                    # 提取首要污染物
                    primary_pollutant = "无"
                    if qaqi_index.get("primaryPollutant"):
                        primary_pollutant = qaqi_index["primaryPollutant"].get(
                            "name", "未知"
                        )

                    # 提取各污染物浓度
                    pollutants = {}
                    for pollutant in hour_data.get("pollutants", []):
                        code = pollutant["code"]
                        pollutants[code] = {
                            "value": pollutant["concentration"]["value"],
                            "unit": pollutant["concentration"]["unit"],
                        }

                    hourly_forecast.append(
                        {
                            "time": forecast_time,
                            "aqi": qaqi_index["aqi"],
                            "aqiDisplay": qaqi_index["aqiDisplay"],
                            "category": qaqi_index["category"],
                            "primaryPollutant": primary_pollutant,
                            "pollutants": pollutants,
                        }
                    )

                return format_air_quality(
                    {
                        "city": city_info["name"],
                        "type": "hourly",
                        "hours": hourly_forecast,
                    }
                )

            elif forecast_type == "daily":
                if "days" not in data:
                    raise ValueError("每日预报API错误：未返回数据")

                # 处理每日预报数据
                daily_forecast = []
                for day_data in data["days"]:
                    # 提取日期（从forecastStartTime中取前10位）
                    forecast_date = day_data["forecastStartTime"][:10]

                    # 查找首选指数 (QAQI)
                    qaqi_index = next(
                        (idx for idx in day_data["indexes"] if idx["code"] == "qaqi"),
                        None,
                    )

                    if not qaqi_index:
                        qaqi_index = day_data["indexes"][0]

                    # 提取首要污染物
                    primary_pollutant = "无"
                    if qaqi_index.get("primaryPollutant"):
                        primary_pollutant = qaqi_index["primaryPollutant"].get(
                            "name", "未知"
                        )

                    # 提取各污染物浓度
                    pollutants = {}
                    for pollutant in day_data.get("pollutants", []):
                        code = pollutant["code"]
                        pollutants[code] = {
                            "name": pollutant["name"],
                            "fullName": pollutant["fullName"],
                            "value": pollutant["concentration"]["value"],
                            "unit": pollutant["concentration"]["unit"],
                        }

                    daily_forecast.append(
                        {
                            "date": forecast_date,
                            "aqi": qaqi_index["aqi"],
                            "aqiDisplay": qaqi_index["aqiDisplay"],
                            "category": qaqi_index["category"],
                            "primaryPollutant": primary_pollutant,
                            "pollutants": pollutants,
                        }
                    )

                return format_air_quality(
                    {"city": city_info["name"], "type": "daily", "days": daily_forecast}
                )

        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


def format_air_quality(air_quality_data: Dict) -> str:
    """
    格式化空气质量数据（支持实时和小时预报）

    参数:
        air_quality_data: 空气质量原始数据，包含type字段标识数据类型

    返回:
        格式化后的空气质量字符串
    """
    if air_quality_data["type"] == "realtime":
        # 实时空气质量数据格式化
        aqi = air_quality_data["aqi"]
        category = air_quality_data["category"]
        primary_pollutant_name = air_quality_data["primaryPollutant"]["name"]
        index_name = air_quality_data["index_name"]

        formatted = f"{air_quality_data['city']}实时空气质量:\n"
        formatted += f"空气质量指数({index_name}): {aqi} ({category})\n"
        formatted += f"首要污染物: {primary_pollutant_name}\n"

        # 添加颜色信息（如果存在）
        if air_quality_data["color"]:
            color = air_quality_data["color"]
            formatted += (
                f"颜色标识: RGB({color['red']}, {color['green']}, {color['blue']})\n"
            )

        # 添加健康建议（如果存在）
        if air_quality_data["health_effect"]:
            formatted += f"健康影响: {air_quality_data['health_effect']}\n"

        if air_quality_data["health_advice_general"]:
            formatted += f"一般人群建议: {air_quality_data['health_advice_general']}\n"

        if air_quality_data["health_advice_sensitive"]:
            formatted += (
                f"敏感人群建议: {air_quality_data['health_advice_sensitive']}\n"
            )

        formatted += "各污染物浓度:\n"
        for pollutant in air_quality_data["pollutants"].values():
            formatted += (
                f"- {pollutant['name']}: {pollutant['value']} {pollutant['unit']}\n"
            )

        return formatted

    elif air_quality_data["type"] == "hourly":
        # 小时预报数据格式化
        formatted = f"{air_quality_data['city']}小时空气质量预报:\n"
        formatted += "时间 | AQI | 等级 | 首要污染物 | 主要污染物浓度\n"
        formatted += "-" * 60 + "\n"

        for hour in air_quality_data["hours"]:
            # 构建污染物浓度字符串
            pollutants_str = ", ".join(
                [
                    f"{p['name']}: {p['concentration']['value']}{p['concentration']['unit']}"
                    for p in hour["pollutants"].values()
                ]
            )

            formatted += (
                f"{hour['time']} | {hour['aqiDisplay']} | {hour['category']} | "
                f"{hour['primaryPollutant']} | {pollutants_str}\n"
            )

        return formatted

    elif air_quality_data["type"] == "daily":
        # 每日预报数据格式化
        formatted = f"{air_quality_data['city']}每日空气质量预报:\n"
        formatted += "日期 | AQI | 等级 | 首要污染物 | 主要污染物浓度\n"
        formatted += "-" * 60 + "\n"

        for day in air_quality_data["days"]:
            # 构建污染物浓度字符串
            pollutants_str = ", ".join(
                [
                    f"{p['name']}: {p['value']}{p['unit']}"
                    for p in day["pollutants"].values()
                ]
            )

            formatted += (
                f"{day['date']} | {day['aqiDisplay']} | {day['category']} | "
                f"{day['primaryPollutant']} | {pollutants_str}\n"
            )

        return formatted

    else:
        raise ValueError("不支持的空气质量数据类型")


@mcp.tool()
async def get_weather_indices(city: str, days: str = "1d") -> str:
    """
    获取城市天气生活指数（格式化输出）

    参数:
        city: 城市名称，例如："北京"、"上海"
        days: 天数，可选值："1d" - 1天生活指数，"3d" - 3天生活指数

    返回:
        格式化后的天气生活指数字符串

    异常:
        ValueError: 参数错误或API响应格式错误
        httpx.RequestError: 网络请求失败
    """
    # 更新为新的API端点格式
    HEFENG_INDICES_URL = "https://nb6yvwk53a.re.qweatherapi.com/v7/indices/{days}"

    valid_days = ["1d", "3d"]
    if days not in valid_days:
        raise ValueError(f"Invalid days parameter. Must be one of: {valid_days}")

    location_info = await get_city_coordinates(city)
    location_id = location_info["id"]
    api_key = get_api_key()
    headers = {"X-QW-Api-Key": api_key}
    params = {
        "location": location_id,
        "type": "0",  # 获取所有生活指数类型
    }

    async with httpx.AsyncClient() as client:
        try:
            response = await client.get(
                HEFENG_INDICES_URL.format(days=days),
                headers=headers,
                params=params,
            )
            response.raise_for_status()
            data = response.json()

            if data.get("code") != "200":
                raise ValueError(
                    f"Weather Indices API Error: {data.get('message', 'Unknown error')}"
                )
            return format_indices(data)
        except httpx.RequestError as e:
            raise httpx.RequestError(f"请求失败: {str(e)}")


def format_indices(indices_data: Dict) -> str:
    """
    格式化天气指数数据（支持按日期分组）

    参数:
        indices_data: 天气指数原始数据

    返回:
        格式化后的天气指数字符串
    """
    if "daily" not in indices_data or not indices_data["daily"]:
        return "无天气指数数据"

    # 按日期分组指数
    indices_by_date = {}
    for index in indices_data["daily"]:
        date = index.get("date", "未知日期")
        if date not in indices_by_date:
            indices_by_date[date] = []
        indices_by_date[date].append(index)

    # 构建格式化输出
    formatted = ""
    for date, indices in indices_by_date.items():
        formatted += f"【{date}】天气指数预报:\n"
        for idx in indices:
            # 处理可能缺失的字段
            name = idx.get("name", "未知指数")
            level = idx.get("level", "-")
            category = idx.get("category", "未知等级")
            text = idx.get("text", "暂无详细描述")

            formatted += f"- {name}: 等级{level} ({category})\n" f"  描述: {text}\n"
        formatted += "\n"

    return formatted


def main():
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
